import os

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from pydantic.v1 import BaseModel, Field

# 读取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")

# 创建大语言模型实例
model = ChatOpenAI(
    model="qwen-plus-latest",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=SecretStr(api_key),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)


class Classification(BaseModel):
    """

    """
    sentiment: str = Field(..., enum=['happy', 'neutral', 'sad'], description="文本的情感")
    aggressiveness: int = Field(..., enum=[1, 2, 3, 4, 5], description="描述文本的攻击性，数字越大表示攻击性越大")
    language: str = Field(...,enum = ['spanish','english','french','中文'],description="文本使用的语言")


tagging_prompt = ChatPromptTemplate.from_template(
    """
    从以下段落中获取所需信息
    只提取'Classification'类中提到的树形。
    段落：{input}
    """
)

chain = tagging_prompt | model.with_structured_output(Classification)

input_text = "中国人民大学的王教授：师德败坏，做出的事情实在让我生气！"

result = Classification = chain.invoke({'input': input_text})

print(result)
